# 解决OpenMP库冲突问题
import os

import agent

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

import numpy as np
import torch
import math

from agent import Agent


class TargetProbabilityMap:
    def __init__(self, grid_resolution: int, target_id, width: int, height: int, initial_prob=None,
                 max_speed: float = agent.AGENT_PARAMS['BSV_MAX_VEL'],
                 time: float = 0.0, perdict_interval=120):
        """
        单个目标的概率地图

        参数:
        grid_resolution: 网格尺寸/分辨率(m)
        target_id: 目标唯一标识符
        width, height: 地图的宽度和高度（网格单元）
        initial_prob: 初始概率分布
        kernel_size: 扩散核大小
        sigma: 高斯扩散核标准差
        """
        self.grid_resolution = grid_resolution
        self.target_id = target_id
        self.width = width
        self.height = height
        self.belief = 1.0  # 目标存在概率
        self.perdict_interval = perdict_interval
        # 初始化概率地图
        if initial_prob is not None:
            assert initial_prob.shape == (width, height)
            self.prob_map = initial_prob.copy()
            # 找到概率最高的位置
            max_idx = np.unravel_index(np.argmax(self.prob_map), self.prob_map.shape)
            grid_x, grid_y = max_idx
            # 将索引转换为二维坐标
            self.position = [grid_x * self.grid_resolution + 0.5 * self.grid_resolution,
                             grid_y * self.grid_resolution + 0.5 * self.grid_resolution]

            # 使用概率图作为权重计算加权平均位置
            x_indices, y_indices = np.meshgrid(np.arange(self.width), np.arange(self.height), indexing='ij')
            weighted_x = np.sum(x_indices * self.prob_map)
            weighted_y = np.sum(y_indices * self.prob_map)
            self.position = [weighted_x * self.grid_resolution + 0.5 * self.grid_resolution,
                             weighted_y * self.grid_resolution + 0.5 * self.grid_resolution]
        else:
            self.prob_map = np.ones((width, height)) / (width * height)
            # 位置信息
            self.position = [width * grid_resolution / 2, height * grid_resolution / 2]

        self.last_detect_position = self.position

        # 速度信息
        self.velocity = [0.0, 0.0]
        self.max_speed = max_speed
        # 最后发现时间
        self.last_detected = -1
        self.time = time

        # 计算卷积核尺寸
        self.kernel_size = 1 + math.ceil(1 + 4 * self.max_speed * self.perdict_interval / self.grid_resolution)
        if self.kernel_size == 1:
            self.kernel_size = 3  # 卷积核最小尺寸为3*3
        # 确保核大小为奇数
        if self.kernel_size % 2 == 0:
            self.kernel_size += 1

        # 创建默认扩散核
        # self.default_kernel = self._create_gaussian_kernel(3 * self.max_speed * self.perdict_interval / grid_resolution)
        self.default_kernel = self.create_velocity_kernel(self.max_speed * self.perdict_interval / self.grid_resolution,
                                                          math.pi,
                                                          self.max_speed * self.perdict_interval / self.grid_resolution / 3,
                                                          math.pi / 12)
        self.current_kernel = self.default_kernel.copy()

        # 初始化信息熵
        self.entropy = self.compute_entropy()  # 计算初始信息熵

    def _create_gaussian_kernel(self, sigma):
        """创建高斯扩散核"""
        ax = np.linspace(-(self.kernel_size - 1) / 2., (self.kernel_size - 1) / 2., self.kernel_size)
        xx, yy = np.meshgrid(ax, ax)
        kernel = np.exp(-(xx ** 2 + yy ** 2) / (2. * sigma ** 2))
        return kernel / kernel.sum()

    def create_velocity_kernel(self, v_magnitude, v_direction, sigma_v, sigma_theta):
        """
        创建基于速度的运动模型核

        参数:
        v_magnitude: 速度大小（网格单位/时间步）
        v_direction: 速度方向（弧度）
        sigma_v: 速度大小的标准差
        sigma_theta: 方向角的标准差（弧度）
        kernel_size: 卷积核大小（奇数），如果为None则自动计算
        """
        # 创建网格
        half_size = self.kernel_size // 2
        y, x = np.mgrid[-half_size:half_size + 1, -half_size:half_size + 1]

        # 计算每个位置相对于速度方向的位移
        displacement = np.sqrt(x ** 2 + y ** 2)
        angle = np.arctan2(-y, -x)  # 计算每个位置的角度

        # 计算方向差异（考虑角度周期性）
        angle_diff = angle - v_direction
        angle_diff = np.arctan2(np.sin(angle_diff), np.cos(angle_diff))
        angle_diff[half_size, half_size] = 0.1  # 中心节点视作总是与速度方向一致，添加微小偏移以避免中心节点概率占比过大
        # 计算方向概率（高斯分布）
        direction_prob = np.exp(-angle_diff ** 2 / (2 * sigma_theta ** 2))
        direction_prob /= direction_prob.sum()

        # 计算速度大小概率（高斯分布）
        magnitude_prob = np.exp(-(displacement - v_magnitude) ** 2 / (2 * sigma_v ** 2))
        magnitude_prob /= magnitude_prob.sum()
        # 组合概率（方向概率和大小概率的乘积）
        kernel = direction_prob * magnitude_prob

        # 归一化
        if kernel.sum() > 0:
            kernel /= kernel.sum()
        else:
            # 若概率全为0，构造一个中心为1的核
            kernel = np.zeros_like(kernel)
            kernel[half_size, half_size] = 1.0
        return kernel.transpose()

    def compute_entropy(self):
        """计算当前概率分布的信息熵"""
        # 展平概率矩阵并移除零概率元素
        if self.belief == 0:  # 不存在的目标熵值视为0
            return 0
        probs = self.prob_map.flatten()
        non_zero_probs = probs[probs > 0]

        # 计算信息熵: H = -sum(p * log2(p)) 并归一化
        entropy = -np.sum(non_zero_probs * np.log2(non_zero_probs)) / math.log2(self.width * self.height)
        self.entropy = entropy
        return entropy

    def get_entropy(self):
        """获取当前信息熵值"""
        return self.entropy

    def predict(self):
        """预测步：使用反射填充进行卷积，避免边界概率异常"""

        # 根据目标状态选择扩散核
        if self.last_detected == self.time or self.belief <= 0:  # 位于视野内不扩散
            return
        elif self.last_detected > 0 and self.velocity is not None:  # 曾被发现过
            # 计算经过时间
            undetected_duration = self.time - self.last_detected
            # 计算速度大小和方向
            speed = math.sqrt(self.velocity[0] ** 2 + self.velocity[1] ** 2)
            direction = math.atan2(self.velocity[1], self.velocity[0])

            # undetected_duration为0时t为0，undetected_duration为无穷时t为1
            # TODO: 参数化调节不确定性提升的速度，以及最大不确定性
            t_v = math.tanh(0.001 * undetected_duration)
            t_a = math.tanh(0.0002 * undetected_duration)

            # 速度大小不确定性：随时间增加
            sigma_v = min(1 + (t_v / 2 + 0.5) * speed, self.max_speed) * self.perdict_interval / self.grid_resolution / 3
            # 方向不确定性：随时间增加
            sigma_theta = min(math.pi / 12, 0.1 + t_a * math.pi / 18)  # 最大10度不确定性

            # 创建速度感知的扩散核
            grid_speed = speed * self.perdict_interval / self.grid_resolution
            self.current_kernel = self.create_velocity_kernel(
                grid_speed, direction, sigma_v, sigma_theta
            )
        elif self.current_kernel is not self.default_kernel:
            self.current_kernel = self.default_kernel
        else:
            pass  # 未曾被发现过的目标使用默认全向高斯核

        # 使用PyTorch进行卷积
        prob_tensor = torch.tensor(self.prob_map, dtype=torch.float32).unsqueeze(0).unsqueeze(0)
        kernel_tensor = torch.tensor(self.current_kernel, dtype=torch.float32).unsqueeze(0).unsqueeze(0)

        # 计算填充大小
        kernel_size = self.current_kernel.shape[0]
        padding = kernel_size // 2

        # 使用反射填充（边界镜像）
        padded_prob = torch.nn.functional.pad(
            prob_tensor,
            (padding, padding, padding, padding),
            mode='reflect'  # 使用反射边界，避免0填充导致边缘部分概率不正常削减
        )

        # 卷积操作（不再需要额外填充）
        diffused = torch.nn.functional.conv2d(
            padded_prob,
            kernel_tensor,
            padding=0
        ).squeeze().numpy()

        # 归一化并更新概率地图
        if diffused.sum() == 0:
            self.reset_to_uniform()
        else:
            self.prob_map = diffused / diffused.sum()
            nonzero_indices = np.nonzero(self.prob_map > 0)
            if len(nonzero_indices[0]) > 0:  # 确保有非零概率点
                x_indices, y_indices = nonzero_indices
                weighted_x = np.sum(x_indices * self.prob_map[x_indices, y_indices])
                weighted_y = np.sum(y_indices * self.prob_map[x_indices, y_indices])
                self.position = [weighted_x * self.grid_resolution + 0.5 * self.grid_resolution,
                                 weighted_y * self.grid_resolution + 0.5 * self.grid_resolution]

        # 信息熵在每步更新后再更新，以降低计算量

    def update_position(self, position, velocity, time: float):
        """
        根据探测器观测更新目标位置

        参数:
        position: 目标位置 (x, y) 实际坐标
        velocity: 目标速度 (vx, vy)
        time: 当前时间
        """
        if position is None:
            raise ValueError("Position cannot be None")

        # 更新位置与速度
        self.position = position
        self.last_detect_position = position
        self.velocity = velocity
        self.last_detected = time
        self.time = time

        # 转换为网格坐标
        grid_x, grid_y = self._real_to_grid(position)

        # 检查边界有效性
        if 0 <= grid_x < self.width and 0 <= grid_y < self.height:
            # 创建新概率图
            self.prob_map = np.zeros((self.width, self.height))
            self.prob_map[grid_x, grid_y] = 1.0
            self.belief = 1.0  # 重置存在概率
        else:
            # 位置超出边界时的处理
            # 保持当前概率图不变，但降低存在概率
            self.belief *= 0.9
            if self.belief < 0.01:
                self.reset_to_uniform()

        # 信息熵为0（完全确定位置）
        self.entropy = 0.0

    def update_absence(self, detector: 'Agent', time: float):
        """
        根据未探测到目标更新地图，支持圆形或扇形探测范围

        参数:
        detector: 尝试进行探测的智能体
        """
        if self.belief == 0: return  # 目标确定不存在的概率地图不更新
        # 更新时间
        self.time = time
        # 预计算常用值，避免重复计算
        grid_res_half = self.grid_resolution * 0.5
        sensor_range_extended = detector.sensor_range + self.grid_resolution * 0.8
        sensor_range_sq = sensor_range_extended ** 2

        # 只计算探测器范围内的网格，避免全图计算
        # 计算探测器网格位置
        detector_grid_x, detector_grid_y = self._real_to_grid((detector.x, detector.y))

        # 计算探测器范围覆盖的网格半径
        grid_radius = int(np.ceil(sensor_range_extended / self.grid_resolution))

        # 计算需要检查的网格范围
        x_min = max(0, detector_grid_x - grid_radius)
        x_max = min(self.width, detector_grid_x + grid_radius + 1)
        y_min = max(0, detector_grid_y - grid_radius)
        y_max = min(self.height, detector_grid_y + grid_radius + 1)

        # 如果没有需要检查的范围，直接返回
        if x_min >= x_max or y_min >= y_max:
            return

        # 只处理探测器范围内的网格
        # 创建局部视图 (避免复制整个数组)
        local_view = self.prob_map[x_min:x_max, y_min:y_max]

        # 计算局部网格中心坐标
        x_indices = np.arange(x_min, x_max)
        y_indices = np.arange(y_min, y_max)

        # 使用广播计算坐标差
        center_x = (x_indices + 0.5) * self.grid_resolution - detector.x
        center_y = (y_indices + 0.5) * self.grid_resolution - detector.y

        # 计算距离平方 (避免开方)
        dx = center_x[:, np.newaxis]  # 形状变为 (x_range, 1)
        dy = center_y[np.newaxis, :]  # 形状变为 (1, y_range)
        distances_sq = dx ** 2 + dy ** 2

        if detector.sensor_angle >= math.pi:  # 全向探测器（圆形）
            mask = (distances_sq <= sensor_range_sq)
        else:  # 定向探测器（扇形）
            # 计算探测器当前朝向
            if detector.vx == 0 and detector.vy == 0:
                current_heading = detector.target_heading
            else:
                current_heading = math.atan2(detector.vy, detector.vx)

            # 计算角度（从探测器指向网格中心）
            angles = np.arctan2(dy, dx)

            # 计算角度差（考虑360度循环）
            angle_diff = np.abs(angles - current_heading)
            angle_diff = np.minimum(angle_diff, 2 * np.pi - angle_diff)

            mask = (distances_sq <= sensor_range_sq) & (angle_diff <= detector.sensor_angle)

        # 6. 只在探测范围内有概率时才更新
        if np.any(local_view[mask] > 0):
            # 计算被清除的总概率
            cleared_prob = np.sum(local_view[mask])
            local_view[mask] = 0

            # 更新整个概率图的总和
            total_prob = np.sum(self.prob_map)  # 比 self.prob_map.sum() 更高效

            if total_prob > 0:
                if cleared_prob > 0:
                    # 归一化处理 - 只计算一次
                    self.prob_map /= total_prob

                    # 使用最大概率点作为位置
                    max_idx = np.unravel_index(np.argmax(self.prob_map), self.prob_map.shape)
                    grid_x, grid_y = max_idx
                    self.position = [
                        grid_x * self.grid_resolution + grid_res_half,
                        grid_y * self.grid_resolution + grid_res_half
                    ]
                    # nonzero_indices = np.nonzero(self.prob_map > 0)
                    # if len(nonzero_indices[0]) > 0:  # 确保有非零概率点
                    #     x_indices, y_indices = nonzero_indices
                    #     weighted_x = np.sum(x_indices * self.prob_map[x_indices, y_indices])
                    #     weighted_y = np.sum(y_indices * self.prob_map[x_indices, y_indices])
                    #     self.position = [weighted_x * self.grid_resolution + grid_res_half,
                    #                      weighted_y * self.grid_resolution + grid_res_half]


            elif self.last_detected > 0:
                # 如果所有概率都为零，则重置为圆内均匀分布
                undetected_duration = self.time - self.last_detected
                radius = self.max_speed * undetected_duration
                self.belief *= 0.7 # 降低存在概率
                # 计算圆内网格数量
                grid_radius = int(np.ceil(radius / self.grid_resolution)) + 1
                center_x, center_y = self._real_to_grid(self.last_detect_position)

                # 确定圆内网格范围
                x_min = max(0, center_x - grid_radius)
                x_max = min(self.width, center_x + grid_radius + 1)
                y_min = max(0, center_y - grid_radius)
                y_max = min(self.height, center_y + grid_radius + 1)

                # 创建新概率图
                self.prob_map = np.zeros((self.width, self.height))
                count = 0

                # 遍历圆内网格
                for i in range(x_min, x_max):
                    for j in range(y_min, y_max):
                        # 计算网格中心到圆心的距离
                        grid_center_x = (i + 0.5) * self.grid_resolution
                        grid_center_y = (j + 0.5) * self.grid_resolution
                        dx = grid_center_x - self.last_detect_position[0]
                        dy = grid_center_y - self.last_detect_position[1]
                        distance = math.sqrt(dx ** 2 + dy ** 2)

                        if distance <= grid_radius * self.grid_resolution:
                            self.prob_map[i, j] = 1.0
                            count += 1
                # 归一化
                self.prob_map /= count
            else:
                # 如果圆内没有网格，重置为全局均匀分布
                self.reset_to_uniform()
                # 使用最大概率点作为位置
                max_idx = np.unravel_index(np.argmax(self.prob_map), self.prob_map.shape)
                grid_x, grid_y = max_idx
                self.position = [
                    grid_x * self.grid_resolution + grid_res_half,
                    grid_y * self.grid_resolution + grid_res_half
                ]

    def reset_to_uniform(self):
        """重置为均匀分布"""
        self.prob_map = np.ones((self.width, self.height)) / (self.width * self.height)
        self.belief = 0.01

    def _real_to_grid(self, real_pos):
        """实际坐标转网格坐标"""
        x, y = real_pos
        grid_x = int(x / self.grid_resolution)
        grid_y = int(y / self.grid_resolution)
        return grid_x, grid_y
